﻿using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Security.Cryptography;
using System.Text;
using NSsh.Common;
using NSsh.Common.Packets;
using NSsh.Common.Types;
using NSsh.Common.Utility;
using NSsh.Server.Services;
using NSsh.Server.TransportLayer;
using NSsh.Server.TransportLayer.State;
using Mono.Math;

namespace NSsh.Client {

    public class ClientVersionsExchangedState : AbstractTransportState {

        public static readonly IList<PublicKeyAlgorithm> preferredHostKeyAlorithms =
            new List<PublicKeyAlgorithm> { PublicKeyAlgorithm.DSA }; //, PublicKeyAlgorithm.RSA 

        public static readonly IList<EncryptionAlgorithm> preferredEncryptionAlorithms =
            new List<EncryptionAlgorithm>
            {
                EncryptionAlgorithm.Aes128Cbc,
                EncryptionAlgorithm.BlowfishCbc,
                EncryptionAlgorithm.TripleDesCbc
            };

        public static readonly IList<MacAlgorithm> preferredMacAlgorithms =
            new List<MacAlgorithm> { MacAlgorithm.HmacSha1 };

        public static readonly IList<CompressionAlgorithm> preferredCompressionAlgorithms =
            new List<CompressionAlgorithm> { CompressionAlgorithm.None };

        private bool _dhReplyPacketReceived;
        private bool _newKeysPacketReceived;

        public override void Process(ITransportLayerManager manager) {

            sendKexInitPacket(manager);

            while (manager.Parameters == null) {
                ProcessPacket(manager);
            }

            SendKexDHInitPacket(manager);

            while (!_dhReplyPacketReceived) {
                ProcessPacket(manager);
                manager.WritePacket(new NewKeysPacket());
            }

            while (!_newKeysPacketReceived) {
                ProcessPacket(manager);
            }

            manager.ChangeState(TransportLayerState.KeysExchanged);
        }

        private void sendKexInitPacket(ITransportLayerManager manager) {
            KexInitPacket kexInit = new KexInitPacket();
            kexInit.KexAlgorithms.Names.Add("diffie-hellman-group1-sha1");

            foreach (PublicKeyAlgorithm algorithm in preferredHostKeyAlorithms) {
                kexInit.ServerHostKeyAlgorithms.Names.Add(PublicKeyAlgorithmHelper.ToString(algorithm));
            }

            foreach (EncryptionAlgorithm algorithm in preferredEncryptionAlorithms) {
                kexInit.EncryptionAlgorithmsClientToServer.Names.Add(EncryptionAlgorithmHelper.ToString(algorithm));
                kexInit.EncryptionAlgorithmsServerToClient.Names.Add(EncryptionAlgorithmHelper.ToString(algorithm));
            }

            foreach (MacAlgorithm algorithm in preferredMacAlgorithms) {
                kexInit.MacAlgorithmsClientToServer.Names.Add(MacAlgorithmHelper.ToString(algorithm));
                kexInit.MacAlgorithmsServerToClient.Names.Add(MacAlgorithmHelper.ToString(algorithm));
            }

            foreach (CompressionAlgorithm algorithm in preferredCompressionAlgorithms) {
                kexInit.CompressionAlgorithmsClientToServer.Names.Add(CompressionAlgorithmHelper.ToString(algorithm));
                kexInit.CompressionAlgorithmsServerToClient.Names.Add(CompressionAlgorithmHelper.ToString(algorithm));
            }

            manager.ClientKexInit = kexInit;
            manager.WritePacket(kexInit);
        }

        public override void KexInitPacket(ITransportLayerManager manager, KexInitPacket packet) {
            manager.ServerKexInit = packet;

            if (!CheckAlgorithmSupport(packet.KexAlgorithms)) {
                throw new InvalidOperationException("No supported kex init algorithm.");
            }

            TransportLayerParameters parameters = new TransportLayerParameters();

            parameters.HostKeyVerification = DecideHostKeyAlgorithm(packet.ServerHostKeyAlgorithms);

            parameters.ClientToServerEncryption = DecideEncryptionAlgorithm(packet.EncryptionAlgorithmsClientToServer);
            parameters.ServerToClientEncryption = DecideEncryptionAlgorithm(packet.EncryptionAlgorithmsServerToClient);

            parameters.ClientToServerMac = DecideMacAlgorithm(packet.MacAlgorithmsClientToServer);
            parameters.ServerToClientMac = DecideMacAlgorithm(packet.MacAlgorithmsServerToClient);

            parameters.ClientToServerCompression = DecideCompressionAlgorithm(packet.CompressionAlgorithmsClientToServer);
            parameters.ServerToClientCompression = DecideCompressionAlgorithm(packet.CompressionAlgorithmsServerToClient);

            manager.Parameters = parameters;
        }

        public void KexDHReplyPacket(ITransportLayerManager manager, KexDHReplyPacket packet) {

            manager.Key = new BigInteger(packet.F).ModPow(manager.X, SshConstants.DHPrime);

            MemoryStream buffer = new MemoryStream();
            BinaryWriter writer = new BinaryWriter(buffer);

            writer.Write(new SshString(manager.ClientVersion).ToByteArray());
            writer.Write(new SshString(manager.ServerVersion).ToByteArray());
            writer.Write(new SshByteArray((byte)PacketType.KexInit, manager.ClientKexInit.GetPayloadData()).ToByteArray());
            writer.Write(new SshByteArray((byte)PacketType.KexInit, manager.ServerKexInit.GetPayloadData()).ToByteArray());
            writer.Write(packet.HostKeyAndCertificates.ToByteArray());
            writer.Write(manager.E);
            writer.Write(packet.F);
            writer.Write(manager.Key);

            manager.Hash = new SHA1CryptoServiceProvider().ComputeHash(buffer.ToArray());
            manager.SessionId = manager.Hash;
            
            // Check that the hash matches what the server sent.
            bool hashVerified = packet.HostKeyAndCertificates.Verify(manager.Hash, packet.SignatureOfH);
            if (!hashVerified) {
                throw new ApplicationException("The server's hash signature was invalid.");
            }

            _dhReplyPacketReceived = true;
        }

        public override void NewKeysPacket(ITransportLayerManager manager, NewKeysPacket packet) {
            ICipherFactory cipherFactory = Dependency.Resolve<ICipherFactory>();

            SymmetricAlgorithm transmitCipher = cipherFactory.CreateCipher(
                manager.Parameters.ClientToServerEncryption,
                manager.Key,
                manager.Hash,
                manager.SessionId,
                'A',
                'C');

            SymmetricAlgorithm receiveCipher = cipherFactory.CreateCipher(
                manager.Parameters.ServerToClientEncryption,
                manager.Key,
                manager.Hash,
                manager.SessionId,
                'B',
                'D');

            IMacFactory macFactory = Dependency.Resolve<IMacFactory>();

            HashAlgorithm transmitMac = macFactory.CreateMac(
                manager.Parameters.ClientToServerMac,
                manager.Key,
                manager.Hash,
                manager.SessionId,
                'E');

            HashAlgorithm receiveMac = macFactory.CreateMac(
                manager.Parameters.ServerToClientMac,
                manager.Key,
                manager.Hash,
                manager.SessionId,
                'F');

            lock (manager.CommunicationLock) {
                manager.TransmitCipher = transmitCipher;
                manager.ReceiveCipher = receiveCipher;
                manager.TransmitMac = transmitMac;
                manager.ReceiveMac = receiveMac;
            }

            _newKeysPacketReceived = true;
        }

        private void SendKexDHInitPacket(ITransportLayerManager manager) {

            KexDHInitPacket kexDHInit = new KexDHInitPacket();
            manager.E = kexDHInit.E;
            manager.X = kexDHInit.X;
            manager.WritePacket(kexDHInit);
        }

        public bool CheckAlgorithmSupport(NameList kexAlgorithms) {
            return kexAlgorithms.Names.Contains("diffie-hellman-group1-sha1");
        }

        public PublicKeyAlgorithm DecideHostKeyAlgorithm(NameList kexAlgorithms) {
            foreach (PublicKeyAlgorithm algorithm in preferredHostKeyAlorithms) {
                if (kexAlgorithms.Names.Contains(PublicKeyAlgorithmHelper.ToString(algorithm))) {
                    return algorithm;
                }
            }

            throw new InvalidOperationException("No supported host key algorithm.");
        }

        private EncryptionAlgorithm DecideEncryptionAlgorithm(NameList encryptionAlgorithms) {
            foreach (EncryptionAlgorithm algorithm in preferredEncryptionAlorithms) {
                if (encryptionAlgorithms.Names.Contains(EncryptionAlgorithmHelper.ToString(algorithm))) {
                    return algorithm;
                }
            }

            throw new InvalidOperationException("No supported encryption algorithm.");
        }

        private MacAlgorithm DecideMacAlgorithm(NameList macAlgorithms) {
            foreach (MacAlgorithm algorithm in preferredMacAlgorithms) {
                if (macAlgorithms.Names.Contains(MacAlgorithmHelper.ToString(algorithm))) {
                    return algorithm;
                }
            }

            throw new InvalidOperationException("No supported mac algorithm.");
        }

        private CompressionAlgorithm DecideCompressionAlgorithm(NameList compressionAlgorithms) {
            foreach (CompressionAlgorithm algorithm in preferredCompressionAlgorithms) {
                if (compressionAlgorithms.Names.Contains(CompressionAlgorithmHelper.ToString(algorithm))) {
                    return algorithm;
                }
            }

            throw new InvalidOperationException("No supported compression algorithm.");
        }
    }
}
